#https://blog.csdn.net/zsfcg/article/details/112510577


import numpy as np
def compute_optimal_transport(M, r, c, lam, epsilon=1e-8):
    """
    Computes the optimal transport matrix and Slinkhorn distance using the
    Sinkhorn-Knopp algorithm
    Inputs:
        - M : cost matrix (n x m)
        - r : vector of marginals (n, ) 每个人拿的甜点的份数
        - c : vector of marginals (m, ) 每种甜点的数量（份数？）
        - lam : strength of the entropic regularization
        - epsilon : convergence parameter
    Outputs:
        - P : optimal transport matrix (n x m)
        - dist : Sinkhorn distance
            U（r，c)包含为我的同事划分甜点的所有方法
    对于一个矩阵，先逐行做归一化，就是将第一行的每个元素除以第一行每个元素的和，得到新的"第一行";
    对于每行都做相同的操作； 再逐列做归一化。 重复以上的两步，最终可以收敛到一个行和为1，列和也为1的Doubly stochastic matrix。
    """
    n, m = np.shape(M)
    P = np.exp(np.multiply(M,-lam))
    temp11=P.sum()
    P /= temp11
    u = np.zeros(n)
    # normalize this matrix
    temp1=P.sum(1)
    temp12=u - temp1
    temp13=np.abs(temp12)
    A=np.max(temp13)
    while A > epsilon:
        u = P.sum(1)
        P *= (r / u).reshape((-1, 1))#行归r化，注意python中*号含义
        P *= (c / P.sum(0)).reshape((1, -1))#列归c化
        A = np.max(np.abs(u - P.sum(1)))
    return P, np.sum(P * M)

M=[
    [ 2, 2, 1, 0, 0],
    [ 0,-2,-2,-2, 2],
    [ 1, 2, 2, 2,-1],
    [ 2, 1, 0, 1,-1],
    [0.5,2, 2, 1, 0],
    [ 0, 1, 1, 1,-1],
    [-2, 2, 2, 1, 1],
    [ 2, 1, 2, 1,-1],
]
c=[0.2,0.1,0.3,0.2,0.2]#[4,2,6,4,4]/sum([4,2,6,4,4])
r=[0.15,0.15,0.15,0.2, 0.1,  0.1,  0.1,  0.05]#[3,3,3,4,2,2,2,1]/sum([3,3,3,4,2,2,2,1])
lam=5
out1,out2=compute_optimal_transport(M, r, c, lam, epsilon=1e-8)
print('out1:{}'.format(out1),'out2:{}'.format(out2),sep='\n')


# # a=np.array([1,2,3],[4,5,6])
# a=np.array([[1,2,3],[4,5,6]])
# print(a.sum(1))